import numpy as np
from admm.ADMM_update import terminal_args
import coptpy as cp
from coptpy import COPT


def storage_update(
    PEssMax: float,
    eff_ch: float,
    eff_dis: float,
    capacity: float,
    max_changing_N: int,
    cost_operating: float,
    depreciation_per_change: float,
    terminal: terminal_args,
    rou: float,
    T: int,
    interval_seconds: float,
):
    """
    eff_ch 充电效率
    eff_dis 放电效率
    interval_seconds两个时间点之间间隔的秒数
    T必须大于2
    """
    # Create COPT environment
    env = cp.Envr()
    # Create COPT model
    model = env.createModel("device_storage")
    # Add variables
    Pch = model.addVars(
        T,
        lb=0,
        ub=PEssMax,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pch",
    )
    Pdis = model.addVars(
        T,
        lb=0,
        ub=PEssMax,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pch",
    )
    state_ch = model.addVars(
        T,
        lb=0,
        ub=1,
        vtype=COPT.BINARY,
        nameprefix="state_ch",
    )
    change_var = model.addMVar(
        T - 1,
        lb=0,
        ub=1,
        vtype=COPT.BINARY,
        nameprefix="change",
    )
    Ee = model.addVars(
        T,
        lb=0.1 * capacity,
        ub=0.9 * capacity,
        vtype=COPT.CONTINUOUS,
        nameprefix="Ee",
    )

    Pi = model.addVars(
        T,
        lb=-1.0 * COPT.INFINITY,
        ub=COPT.INFINITY,
        vtype=COPT.CONTINUOUS,
        nameprefix="Pch",
    )

    # Add linear constraints
    model.addConstrs(
        [Pch[t] <= state_ch[t] * PEssMax for t in range(0, T)],
        nameprefix="ChargeCons",
    )
    model.addConstrs(
        [Pdis[t] <= (1 - state_ch[t]) * PEssMax for t in range(0, T)],
        nameprefix="DischargeCons",
    )
    model.addConstrs(
        [
            Ee[t]
            == Ee[t - 1]
            + eff_ch * Pch[t - 1] * (interval_seconds / 3600)
            - eff_dis * Pdis[t - 1] / eff_dis * (interval_seconds / 3600)
            for t in range(1, T)
        ],
        nameprefix="EeCons",
    )
    model.addConstr(
        Ee[0]==Ee[T-1],
        name="EeSameC",
    )
    model.addConstrs(
        [
            change_var[t - 1].item() >= (state_ch[t] - state_ch[t - 1])
            for t in range(1, T)
        ],
        nameprefix="change_b1_cons",
    )
    model.addConstrs(
        [
            change_var[t - 1].item() >= (state_ch[t - 1] - state_ch[t])
            for t in range(1, T)
        ],
        nameprefix="change_b2_cons",
    )
    model.addConstr(
        change_var @ np.ones((T - 1, 1)) <= max_changing_N, name="changing_limit"
    )
    model.addConstrs(
        [Pi[t] == Pch[t] - Pdis[t] for t in range(0, T)], nameprefix="terminal_cons"
    )

    # Set linear objective
    obj = cost_operating * (Pch[0] + Pdis[0])
    for t in range(1, T):
        obj = (
            obj
            + cost_operating * (Pch[t] + Pdis[t])
            + depreciation_per_change * change_var[t - 1].item()
        )

    for t in range(0, T):
        obj = obj + terminal.LP[t] * (Pi[t] - terminal.P[t])

    # 二次项
    for t in range(0, T):
        obj = obj + rou / 2 * (Pi[t] - terminal.P[t]) * (Pi[t] - terminal.P[t])

    model.setObjective(obj, COPT.MINIMIZE)

    # Solve the problem
    model.solve()

    # gather solution and update y
    # 返回得到的k+1步交流线路端口状态和第k步的拉格朗日乘子
    if model.status != COPT.OPTIMAL:
        raise RuntimeError("储能更新未能找到最优解")
    else:
        y = {
            "terminal": terminal_args(T, require=("P",)),
            "value": model.value,
            "Pch": np.array([Pch[k].x for k in range(0, T)]),
            "Pdis": np.array([Pdis[k].x for k in range(0, T)]),
            "state_ch": np.array([state_ch[k].x for k in range(0, T)]),
            "Ee": np.array([Ee[k].x for k in range(0, T)]),
        }
        y["terminal"].P = np.array([Pi[k].x for k in range(0, T)]).transpose()
        y["terminal"].LP = terminal.LP
    return y


def start():
    T = 36
    T1 = terminal_args(T,require=('P',))
    T1.P = np.random.random((T, 1))*30-10
    T1.LP = np.random.random((T, 1))+2
    ret=storage_update(50.0,0.98,0.98,5000,10,0.05,4,T1,0.02,T,900)
    print(f'Pk:{T1.P.transpose()}\nPk+1:{ ret["terminal"].P }\nEe:{ret["Ee"]}')
    print(f'is_charge:{ret["state_ch"]}')
